import json
from pathlib import Path
from typing import Literal, Optional, List, Any

import requests
from fastapi import FastAPI, Request, Header
from fastapi import HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from fastapi.staticfiles import StaticFiles
from loguru import logger
from openai import OpenAI
from pydantic import BaseModel, constr

app = FastAPI()
# 确定ui目录的路径
static_dir = Path(__file__).parent.parent / "ui"
# 挂载静态文件
app.mount("/static", StaticFiles(directory=static_dir, html=True), name="static")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 可以是 ["*"] 来允许所有域
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头
)

api_key = ''
base_url = 'https://api.openai.com/v1'
model = ''

client = OpenAI(api_key=api_key, base_url=base_url)

_query_by_contexts_prompt = """
你是一个由周某人开发的AI助手，你叫LIKEGAKKI-BOT。你的任务是在收到用户问题后，提供简洁、准确的回答。你会得到一系列与问题相关的上下文，每个上下文都以类似 [[citation:x]] 这样的引用编号开始，其中 x 是一个数字。如果适用，请使用这些上下文并在句子的末尾进行引用。
你的回答必须是正确的、准确的，并且由专家以中立和专业的语调撰写。不要提供与问题无关的信息，也不要重复内容。如果给定的上下文未能提供足够的信息，请使用“关于...缺少信息”来指出。
请使用引用编号来引用上下文，格式为：[citation:x]，x是你引用的上下文的编号。
如果一句话来源于多个上下文，请列出所有适用的引用，例如：[citation:3][citation:5]。
除了代码和特定名称及引用外，你的回答必须使用与问题相同的语言撰写。

这是一组上下文：

{context}

记住，不要盲目地逐字复述上下文。以下是用户的问题：

{question}
"""

_query_prompt = """
你是一个由周某人开发的AI助手，你叫LIKEGAKKI-BOT。
你的任务是在收到用户问题后，提供简洁、准确的回答。
对于超过你能力范围的问题，请澄清你不知道。
除了代码和特定名称及引用外，你的回答必须使用与问题相同的语言撰写。

用户问题：{question}
"""

_judgment_question_prompt = """
请根据以下提示，判断用户的问题是否属于以下类型之一，这些类型的问题可能超出了你的回答能力：

1. 实时数据查询：如果问题包含“最新”、“当前”、“实时”等关键词，并且需要即时数据，如股票价格、天气等。

2. 特定日期信息：如果问题指定了一个日期，并且需要该日期的特定信息，而这个日期是在我的知识截止日期之后。

3. 个性化推荐：如果问题需要根据用户的个人偏好或历史数据进行推荐。

4. 最新研究：如果问题涉及最新的科学发现或技术进步，并且这些信息是在我的知识截止日期之后。

5. 特定个人问题：如果问题需要了解特定个人当前的状态或偏好。

6. 实时事件进展：如果问题询问的是正在发生或刚刚发生的事件的最新进展。

7. 特定地理位置信息：如果问题需要特定地理位置的详细信息，而这个信息可能不在我的训练数据中。

8. 依赖外部API服务：如果问题需要调用外部服务，如翻译、地图等。

现在，请根据用户的问题内容，使用上述提示来判断问题类型。如果问题属于上述任何一种类型，请直接输出“是”；如果不是，请直接输出“否”。
用户问题：{question}
"""

_guess_your_question = """
你是一个提问专家，根据用户的“原始问题”和“相关上下文”帮助用户提出“相关问题”。请确定值得跟进的主题，并编写不超过20个字的提问。确保包括具体的事件、名称、地点等，使这些问题可以独立提出。记住，根据原始问题和相关上下文，建议三个这样的后续问题。请勿重复原始问题。每个相关问题不超过20个字。例如：

相关上下文：
'''
豆瓣电影【口碑榜】2023-09-11 更新. 影志. 榜单地址：http://www.douban.com/movie/ (页面右侧，每周五下午更新) 口碑榜是豆瓣根据用户最近一周收藏的不低于7分的新片 ...

《动物园里有什么？》 2024-01-12, 喜剧 ; 《大雨》, 2024-01-12, 动画/ 奇幻/ 冒险 ; 《临时劫案》, 2024-01-19, 喜剧/ 动作/ 犯罪 ; 《花千骨》 ...

最新上映的新片,坂本龙一：杰作,哆啦A梦：大雄的地球交响乐,泳者之心,乜代宗师,潜艇总动员：寻龙,加菲猫家族,疯狂的麦克斯：狂暴女神,美国内战,日夜江河,梦想之地, ...

2345电影大全共有10000部电影聚集爱情、动作、喜剧、恐怖、战争、科幻、武侠、冒险、枪战、悬疑、犯罪等2024最新电影。好看的热门电影排行榜推荐，更多电影尽在2345 ...

最近有哪些必看电影推荐呢？2023年必看电影推荐懒人包来了！1月有哪些值得一看的电影推荐呢？许多强档大片纷纷在台湾电影院上映，包含国内台湾 ...

下面小编就为大家盘点，2023年最新十部美国电影，都是在2023年上映的美国大片，如果你想看美国新片，直接参考这份最新美国电影榜单即可。 1、火之谜. 《火 ...

2022年刚过半，电影世界已是妙趣横生。从古灵精怪的《瞬息全宇宙》、描绘动人亲情的《小妈妈》到探讨色情产业的《欢愉》，时报影评人选出了这十部最佳 ...

有什么好看的电影推荐？不能错过的经典电影有哪些？最新热门电影排行榜前十名有什么片？爱情、搞笑、喜剧、文艺、小清新、悬疑、恐怖、惊悚、动画等好看的电影盘点。
'''
原始问题：推荐几部最近好看的电影
相关问题：["豆瓣2024十大电影是什么？","近期的戏剧电影有哪些？","近期有哪些电影不值得看？"]
----------------
相关上下文：
'''
17日（今天）. 晴. 37/24℃. <3级. 18日（明天）. 晴 ...

气温 · 降水 · 风速 ; 32.1℃ · 无降水 · 4.1m/s ; 36.5℃ · 无降水 · 4.4m/s ...

当前天气信息. 天气：晴27℃. 体感：26℃.

北京天气 ; 06-12(今天) · 23～38℃ ; 06-13(明天) · 23～37℃ ; 06-14(周五) · 21～34℃ ; 06-15(周六) · 21～33℃ ; 06-16(周日) · 22～34℃.

北京欢乐谷天气预报，及时准确发布中央气象台天气信息，便捷查询北京欢乐谷今日天气，北京欢乐谷周末天气，北京欢乐谷一周天气预报，北京欢乐谷15日天气预报，北京欢乐 ...

北京, 北京市, 中國Weather Forecast, with current conditions, wind, air quality, and what to expect for the next 3 days.

15天预报 · 周五 晴 晴. 35° 21°. 晴 06/14 · 周六 晴 晴. 32° 21°. 晴 06/15 · 周日 晴 晴. 35° 22°. 晴 06/16 · 周一 晴 晴. 37° 24°. 多云 06/17 · 周二 多云 多云. 37° ...
'''
原始问题：明天北京天气怎么样？
相关问题：["北京周四天气如何？","北京今天风俗多少？","北京明天会下雨吗？"]
----------------
相关上下文：
'''
byte 数据类型是8位、有符号的，以二进制补码表示的整数； · 最小值是-128（-2^7）； · 最大值是127（2^7-1）； · 默认值是0； · byte 类型用在大型数组中节约空间，主要代替 ...

基本数据类型是Java 语言操作数据的基础，包括boolean、char、byte、short、int、long、float 和double，共8 种。 2）引用数据类型。 除了基本数据类型 ...

Java 中的基础数据类型有8 种，分别是：byte, boolean, short, char int, float, long double。 如果要理解这些基础类型如何存储数据，需要知道两个基础知识：.

Java定义了以下几种基本数据类型：. 整数类型：byte，short，int，long. 浮点数类型：float，double. 字符类型：char.

Java基本类型分类、大小及表示范围 · 数值型：. 整数类型：byte、short、int、long; 浮点类型：float、double · 字符型：char · 布尔型：boolean.

基本数据类型--四类八种：整数类（byte、short、int、long）、浮点类（float、double）、字符类（char）、布尔型（boolean）；. 除此之外即为引用类数据 ...

八种基本数据类型的默认值 ; 1. byte(字节). 8-bit ; 2. short(短整数). 16-bit ; 3. int(整数). 32-bit ; 4. long(长整数). 64-bit.

Java基本类型共有八种，基本类型可以分为三类，字符类型char，布尔类型boolean以及数值类型byte、short、int、long、float、double。数值类型又可以分为 ...
'''
原始问题：java基本类型有哪些？
相关问题：["Java中浮点类型大小多少？","Java中short默认值多少？","Java中byte和char的区别是什么？"]
----------------
相关上下文：
'''
{context}
'''
原始问题：{question}
相关问题：
"""


# 调用大模型执行
def completion(task_prompt, history_list: [] = None, contexts: [] = None):
    if not history_list:
        history_list = []

    if contexts:
        final_prompt = _query_by_contexts_prompt.format(
            context="\n\n".join(
                [f"[[citation:{i + 1}]] {c['snippet']}" for i, c in enumerate(contexts)]
            ),
            question=task_prompt
        )
    else:
        final_prompt = _query_prompt.format(question=task_prompt)

    user_message = OpenAIMessage(role='user',
                                 content=final_prompt)

    history_list.append(user_message)

    llm_response_chunks = client.chat.completions.create(
        model=model,
        messages=history_list,
        temperature=0.8,
        top_p=1,
        stream=True
    )

    return llm_response_chunks


REFERENCE_COUNT = 8
SERPER_SEARCH_ENDPOINT = "https://google.serper.dev/search"
SERPER_SEARCH_KEY = ""


def search(query: str, subscription_key: str):
    payload = json.dumps({
        "q": query,
        "num": (
            REFERENCE_COUNT
            if REFERENCE_COUNT % 10 == 0
            else (REFERENCE_COUNT // 10 + 1) * 10
        ),
        "gl": "cn",
        "hl": "zh-cn"
    })
    headers = {"X-API-KEY": subscription_key, "Content-Type": "application/json"}
    logger.info(
        f"{payload} {headers} {subscription_key} {query} {SERPER_SEARCH_ENDPOINT}"
    )
    response = requests.post(
        SERPER_SEARCH_ENDPOINT,
        headers=headers,
        data=payload,
        timeout=5,
    )
    if not response.ok:
        logger.error(f"{response.status_code} {response.text}")
        raise HTTPException(response.status_code, "Search engine error.")
    json_content = response.json()
    contexts = []
    if json_content.get("knowledgeGraph"):
        url = json_content["knowledgeGraph"].get("descriptionUrl") or json_content["knowledgeGraph"].get("website")
        snippet = json_content["knowledgeGraph"].get("description")
        if url and snippet:
            contexts.append({
                "name": json_content["knowledgeGraph"].get("title", ""),
                "url": url,
                "snippet": snippet
            })

    if json_content.get("answerBox"):
        url = json_content["answerBox"].get("link")
        snippet = json_content["answerBox"].get("snippet") or json_content["answerBox"].get("answer")
        if url and snippet:
            contexts.append({
                "name": json_content["answerBox"].get("title", ""),
                "url": url,
                "snippet": snippet
            })

    contexts += [{"name": c["title"], "url": c["link"], "snippet": c.get("snippet", "")} for c in
                 json_content["organic"]]

    return contexts[:REFERENCE_COUNT]


def _stream_response(llm_response_chunks, contexts, more_questions):
    yield json.dumps(contexts)
    yield "\n\n__LLM_RESPONSE__\n\n"
    for chunk in llm_response_chunks:
        if chunk.choices:
            yield chunk.choices[0].delta.content or ""
    if more_questions:
        yield f"\n\n__MORE_QUESTIONS__\n\n{more_questions}"



def _judgment_question(question: str):
    prompt = _judgment_question_prompt.format(question=question)
    messages = [
        {
            "role": "user",
            "content": prompt
        }
    ]

    llm_response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0.1,
        top_p=1
    )

    if llm_response.choices:
        content = llm_response.choices[0].message.content
        if isinstance(content, str):
            return "是" not in content

    return False


def _guess_question(question, contexts):
    messages = [
        {
            "role": "user",
            "content": _guess_your_question.format(
                context="\n\n".join([c["snippet"] for c in contexts]),
                question=question
            )
        }
    ]
    llm_response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0.1,
        top_p=1
    )

    if llm_response.choices:
        content = llm_response.choices[0].message.content
        print(content)
        return content
    return None


class OpenAIMessage(BaseModel):
    role: Literal["user", 'system', "assistant"]
    content: constr(min_length=1)


class ChatMessage(BaseModel):
    question: constr(min_length=1)  # 至少包含一个字符
    history_list: Optional[List[OpenAIMessage]] = None


class ResponseResult(BaseModel):
    code: int
    message: str
    data: Optional[Any] = None

    @classmethod
    def success(cls, message=None, data=None):
        return cls(code=200, message=message if message else "success", data=data)

    @classmethod
    def fail(cls, message=None):
        return cls(code=500, message=message if message else "fail")


@app.post("/api/chat")
async def chat(chat_message: ChatMessage) -> StreamingResponse:
    contexts = []
    # 对问题进行归类，判断当前是否需要进行联网检索
    result = _judgment_question(chat_message.question)
    logger.info(f"当前问题：{chat_message.question}，大模型{'可以' if result else '不能'}回答该问题。")
    more_questions = None

    if not result:
        logger.info("针对问题进行联网检索")
        # 对提出的问题进行检索
        contexts = search(chat_message.question, SERPER_SEARCH_KEY)
        logger.info(f"当前问题：{chat_message.question}，检索到的信息：{contexts}")
        # 猜你想问
        questions = _guess_question(chat_message.question, contexts)
        more_questions = json.loads(questions)
    llm_response_chunks = completion(chat_message.question, chat_message.history_list, contexts)
    return StreamingResponse(_stream_response(llm_response_chunks, contexts, more_questions), media_type="text/html")


if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=9995)
